import pandas as pd
from datetime import datetime
import pytz
from typing import Dict, Tuple
import re

class RestaurantTimeConverter:
    def __init__(self):
        # 创建州名到时区的映射字典
        self.state_timezone_map = {
            'MD': 'America/New_York',     # EDT
            'TX': 'America/Chicago',      # CDT
            'IL': 'America/Chicago',      # CDT
            'CO': 'America/Denver',       # MDT
            'MI': 'America/New_York',     # EDT
            'CA': 'America/Los_Angeles',  # PDT
        }

    def extract_state(self, address: str) -> str:
        """从地址中提取州名缩写"""
        try:
            # 确保地址是字符串类型
            address = str(address).strip()
            # 使用更精确的正则表达式
            state_pattern = r'(?:^|\s)([A-Z]{2})(?:\s|$)'
            match = re.search(state_pattern, address)
            if match:
                state = match.group(1)
                print(f"从地址 '{address}' 提取到州名: {state}")  # 调试信息
                return state
            else:
                print(f"未能从地址 '{address}' 提取到州名")  # 调试信息
                return ''
        except Exception as e:
            print(f"提取州名时出错: {str(e)}")
            return ''

    def get_beijing_time_range(self, local_timezone: str) -> str:
        """计算当地14:00-16:00对应的北京时间"""
        local_tz = pytz.timezone(local_timezone)
        beijing_tz = pytz.timezone('Asia/Shanghai')

        now = datetime.now()
        local_start = local_tz.localize(datetime(now.year, now.month, now.day, 14, 0))
        local_end = local_tz.localize(datetime(now.year, now.month, now.day, 16, 0))

        beijing_start = local_start.astimezone(beijing_tz)
        beijing_end = local_end.astimezone(beijing_tz)

        start_str = beijing_start.strftime('%H:%M')
        end_str = beijing_end.strftime('%H:%M')

        next_day = ''
        if beijing_start.day != local_start.day:
            next_day = '(+1天)'

        return f"{start_str}-{end_str}{next_day}"

    def process_excel(self, input_file: str, output_file: str):
        """处理Excel文件并添加时间转换"""
        try:
            # 读取Excel文件
            df = pd.read_excel(input_file)
            print(f"成功读取Excel文件，共 {len(df)} 行数据")

            # 添加新列
            df['local_time_range'] = '14:00-16:00'
            df['beijing_time_range'] = ''
            df['state'] = ''

            # 处理每一行
            for index, row in df.iterrows():
                address = str(row['restaurant_address'])
                state = self.extract_state(address)
                df.at[index, 'state'] = state
                if state in self.state_timezone_map:
                    timezone = self.state_timezone_map[state]
                    beijing_time = self.get_beijing_time_range(timezone)
                    df.at[index, 'beijing_time_range'] = beijing_time

            # 获取并打印所有唯一的州名
            unique_states = sorted([s for s in df['state'].unique() if s])
            print(f"找到的所有州: {unique_states}")

            # 准备列名
            columns = ['restaurant_name', 'restaurant_address', 'restaurant_Phone',
                      'local_time_range', 'beijing_time_range', 'state']

            new_rows = []
            for state in unique_states:
                # 创建州标题行（使用字典列表而不是Series）
                header_row = dict.fromkeys(columns, '')  # 先创建所有列都为空的行
                header_row['restaurant_name'] = f'=== {state} ==='  # 只设置州名标题
                new_rows.append(header_row)

                # 添加该州的所有餐厅
                state_restaurants = df[df['state'] == state].copy()
                print(f"州 {state} 有 {len(state_restaurants)} 家餐厅")
                state_restaurants = state_restaurants.drop('state', axis=1)
                restaurant_rows = state_restaurants.to_dict('records')
                new_rows.extend(restaurant_rows)

                # 添加空行
                empty_row = dict.fromkeys(columns, '')
                new_rows.append(empty_row)

            # 创建新的DataFrame并指定列顺序
            result_df = pd.DataFrame(new_rows, columns=columns)

            # 保存结果前删除state列
            if 'state' in result_df.columns:
                result_df = result_df.drop('state', axis=1)

            # 保存结果
            result_df.to_excel(output_file, index=False)
            print(f"转换完成，结果已保存至 {output_file}")

        except Exception as e:
            print(f"处理过程中出现错误: {str(e)}")